import os
import pandas as pd
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain_community.embeddings import XinferenceEmbeddings
from tqdm import tqdm
import torch
from typing import List, Dict
from openai import OpenAI

# --------------------------
# 配置部分
# --------------------------
class Config:
    # 嵌入模型配置
    EMBEDDING_MODEL_PATH = "../xinference/modelscope/hub/Xorbits/bge-m3"  # 替换为您的模型路径
    DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
    
    # 文本分割配置
    CHUNK_SIZE = 300
    CHUNK_OVERLAP = 50
    
    # 检索配置
    TOP_K = 5
    SCORE_THRESHOLD = 0.5

    # 查询扩展配置
    QUERY_EXPANSION_TEMPLATE = """
    你是一名政策咨询专家，请根据原始问题生成3个相关的扩展问题。
    原始问题: {question}
    
    要求:
    - 每个扩展问题应该从不同角度探讨原始问题
    - 保持问题简洁明了
    - 用中文输出
    
    请按以下格式返回结果,每个扩展问题占一行，格式为:
    1. 第一个扩展问题
    2. 第二个扩展问题
    3. 第三个扩展问题
    """
    
    # Ollama配置
    # OLLAMA_MODEL = "llama3.2:3b"  # 替换为您本地的模型名称

# --------------------------
# 核心功能实现
# --------------------------
class PolicyRAGSystem:
    def __init__(self):
        # 初始化嵌入模型
        self.embedding_model = XinferenceEmbeddings(
            server_url="http://0.0.0.0:9997", 
            model_uid="bge-m3"
            # model_uid="bge-large-zh"
            # model_uid="bge-large-zh-v1.5"
            # model_uid="Qwen3-Embedding-0.6B"
        )
        # self.embedding_model = HuggingFaceBgeEmbeddings(
        #     model_name=Config.EMBEDDING_MODEL_PATH,
        #     model_kwargs={'device': Config.DEVICE},
        #     encode_kwargs={'normalize_embeddings': False},
        #     query_instruction="为这个政策问题生成表示以用于检索相关文档："
        # )
        self.vector_store = None
        self.client = OpenAI(base_url="http://127.0.0.1:9997/v1", api_key="None")

    def build_from_excel(self, excel_path):
        """从Excel文件构建知识库"""
        # 1. 加载政策数据
        df = pd.read_excel(excel_path)
        
        # 2. 转换为Document对象
        documents = []
        for _, row in tqdm(df.iterrows(), desc="处理政策文件"):
            # 合并关键字段作为内容
            # content = f"""标题：{row['标题']}
            #     发文单位：{row['发文单位']}
            #     文号：{row['文号']}
            #     日期：{row['发布日期']}
            #     正文摘要：{row['正文'][:1000]}..."""  # 限制正文长度
            
            # metadata = {
            #     "source": row['文号'],
            #     "title": row['标题'],
            #     "authority": row['发文单位'],
            #     "date": str(row['发布日期'])
            # }
            content = f"""标题：{row['标题']}
                文号：{row['文号']}
                日期：{row['发文日期']}
                正文摘要：{str(row['正文'])[:500]}..."""  # 限制正文长度
            
            metadata = {
                "source": row['文号'],
                "title": row['标题'],
                "date": str(row['发文日期'])
            }
            documents.append(Document(page_content=content, metadata=metadata))
        
        # 3. 文本分块
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=Config.CHUNK_SIZE,
            chunk_overlap=Config.CHUNK_OVERLAP
        )
        chunks = text_splitter.split_documents(documents)
        
        # 4. 构建向量库
        self.vector_store = FAISS.from_documents(chunks, self.embedding_model)
        print(f"知识库构建完成，共 {len(chunks)} 个文本块")

    def expand_query(self, question: str) -> List[str]:
        """查询扩展 - 生成相关问题"""
        #print("expand_query_IN")
        # response = ollama.chat(
        #     model=Config.OLLAMA_MODEL,
        #     messages=[{"role": "user", "content": Config.QUERY_EXPANSION_TEMPLATE.format(question=question)}]
        # )

        response = self.client.chat.completions.create(
            model="qwen2.5-instruct-int4",
            messages=[
                {"role": "user", "content": Config.QUERY_EXPANSION_TEMPLATE.format(question=question)},
            ],
            temperature=0.2
        )
        #print("相关问题为：", response)
        #print(11111111)
        # print(response['message']['content'])
        #expanded_queries = [q.split(". ")[1] for q in response.choices[0].message.content.split("\n") if q]
        expanded_queries = []
        for line in response.choices[0].message.content.split("\n"):
            # 跳过空行
            if not line.strip():
                continue
                
            # 尝试按序号分割（如 "1. 查询内容"）
            parts = line.split(". ", 1)  # 只分割一次
            
            # 检查分割是否有效（有至少两部分）
            if len(parts) > 1:
                # 取序号后面的内容（索引1）
                query = parts[1].strip()
            else:
                # 如果没有序号，使用整行内容
                query = line.strip()
            
            # 添加到结果列表（确保内容非空）
            if query:
                expanded_queries.append(query)

        #print(22222222)
        #print([question] + expanded_queries[:3])
        return [question] + expanded_queries[:3]  # 原始问题+3个扩展问题

    def retrieve_policies(self, question, top_k=None, threshold=None):
        """检索相关政策"""
        #print("retrieve_policies_IN")
        top_k = top_k or Config.TOP_K
        threshold = threshold or Config.SCORE_THRESHOLD
        
        # 查询扩展
        queries = self.expand_query(question)

        # 带分数的相似度搜索
        # docs_with_scores = self.vector_store.similarity_search_with_score(
        #     question, 
        #     k=top_k*2  # 扩大初选范围
        # )
        # 多查询检索
        all_docs = []
        for query in queries:
            docs_with_scores = self.vector_store.similarity_search_with_score(
                query, 
                k=top_k*2  # 扩大初选范围
            )
            # 过滤低质量结果并去重
            for doc, score in docs_with_scores:
                if score >= threshold and doc.metadata['source'] not in [d.metadata['source'] for d in all_docs]:
                    all_docs.append(doc)
                    if len(all_docs) >= top_k:
                        break
        
        # 过滤低质量结果
        filtered = all_docs[:top_k]
        #print(filtered)
        
        return filtered


        

    def generate_answer(self, question):
        """生成政策解答"""
        #print("generate_answer_IN")
        # 1. 检索相关文档
        docs = self.retrieve_policies(question)
        if not docs:
            print("未找到相关政策信息")
            return "未找到相关政策信息"
        
        #print("已找到相关政策信息：", docs)
        
        # 2. 构建提示词
        context = "\n\n".join(
            f"【政策标题】{doc.metadata['title']}\n"
            f"【政策发文文号】{doc.metadata['source']}\n"
            f"【发布时间】{doc.metadata['date']}\n"
            # f"【发文单位】{doc.metadata['authority']}\n"
            f"【内容摘要】{doc.page_content[:300]}..."
            for doc in docs
        )
        
        # print(context)

        template = f"""
        你是一名政策咨询专家，请根据以下相关政策内容回答问题。
        
        ### 相关政策：
        {context}
        
        ### 问题：
        {question}
        
        ### 要求：
        - 只基于提供的政策信息回答
        - 先给出政策标题（标题内容不要修改，基于原文不变）
        - 标明发文单位和日期，单位需要全部写出，日期精确到日即可，例如“2021年1月1日”
        - 简要介绍政策内容
        
        ### 输出格式如下，请严格按照这个格式，不要有多余的内容，例如'#','*'等符号：
        标题：xxx
        发文单位：xxx
        文号：xxx  
        日期：xxx
        正文摘要：xxx
        
        ### 参考结果，仅作为参考：
        标题：天津市人民政府关于印发天津市推动科技创新和产业创新深度融合工作方案（2025—2027年）的通知
        发文单位： 天津市人民政府
        文号： 津政发〔2025〕9号
        日期： 2025年05月17日
        正文摘要： 该政策是针对推动科技创新和产业创新深度融合的工作方案，其主要目的是加快培育和发展新质生产力，支撑科技强市建设。方案中提到了多个方面的支持与措施，包括研发投入增加、企业研发能力建设等

        ### 注意事项
        - 避免主观臆断
        - 不要修改标题、日期等内容
        - 如果question和政策没有任何的关系，那么请直接返回“请明确表述您的问题”
        - 在给出政策之后，回答问题（如果仅仅是问政策或者问政策内容，那么可以简要介绍政策内容，如果要求详细解释，那么详细解释并总结政策内容），其他情况下，需要保证完整回答问题
        """
        
        # 3. 调用本地Ollama模型
        # response = ollama.chat(
        #     model=Config.OLLAMA_MODEL,
        #     messages=[{"role": "user", "content": template}]
        # )
        # return response['message']['content']
        return self.client.chat.completions.create(
            model="qwen2.5-instruct-int4",
            messages=[
                {"role": "user", "content": template}
            ],
            temperature=0.2
        ).choices[0].message.content
        

    
# --------------------------
# 使用示例
# --------------------------

if __name__ == "__main__":
    # 初始化系统
    policy_rag = PolicyRAGSystem()

    excel = "政策总表.xlsx"
    dir = "total_db"
    
    # 构建知识库（只需运行一次）
    if not os.path.exists(dir):  # 假设没有构建过
        policy_rag.build_from_excel(excel_path=excel)  # 替换为您的Excel路径
        policy_rag.vector_store.save_local(dir)
    else:
        policy_rag.vector_store = FAISS.load_local(
            dir, 
            policy_rag.embedding_model,
            allow_dangerous_deserialization=True
        )
    
    # 交互式查询
    while True:
        question = input("\n请输入您想查询的政策问题（输入q退出）: ")
        if question.lower() == 'q':
            break
            
        print("\n正在检索相关政策...")
        answer = policy_rag.generate_answer(question)
        print("\n政策解答：")
        print(answer)